
"""
获取数据表

"""
from langgraph.types import Command
from core.state import AnalyzeState
from langgraph.config import get_stream_writer
from core.llm import dataset_filter_analyzer
from loguru import logger
from models.aiquery import AiQueryData, TableData

def dataset_filter_node(state: AnalyzeState):

    writer = get_stream_writer()
    writer("对数据进行过滤")
    logger.info("对数据进行过滤")

    datasets = state.get('datasets')

    if len(datasets) == 0:
        filtered_dataset = []
    else:
        filtered_dataset = dataset_filter_analyzer.invoke(
            query=state.get('query'),
            datasets=datasets
        )
    logger.info(f'过滤后的数据集个数：{len(filtered_dataset)}')

    used_queries = set()
    for dataset in filtered_dataset:
        if isinstance(dataset, AiQueryData) and isinstance(dataset.data, TableData):
            used_queries.add(dataset.data.query)

    query = state.get('query')
    if query not in used_queries:
        filtered_dataset.append(AiQueryData(
            description=query,
            data=None
        ))
    writer(filtered_dataset)
    logger.info(f'最终返回的数据集：{filtered_dataset}')

    return Command(update={"filtered": filtered_dataset})